package nmf;

import org.ejml.simple.SimpleMatrix;

public class NMFAlgorithm {
   static double epsilon = .000001;
   static int bestRValue = -1;
   
   public static NMFAlgorithm.Results getRValue(final SparseMatrix V) {
      final NMFAlgorithm.Results[] results = new NMFAlgorithm.Results[3];
      //double M = Math.sqrt(V.num_rows);
      final double M = .025 * V.num_rows;
      final int[] rValues = { (int) M, (int) (M / 2), (int) (3 * M / 2) };
      int lowest = 0;
      
      for (int i = 0; i < 1; i++) {
         results[i] = factor(V, rValues[i]);
      }
      
      for (int i = 1; i < 1; i++) {
         if (results[i].error < results[lowest].error) {
            lowest = i;
         }
      }
      
      bestRValue = rValues[lowest];
      return results[lowest];
   }
   
   public static int getNumIterations(final SparseMatrix V) {
      return 20;
   }
   
   public static double FrobeniusCost(final SparseMatrix V, final SparseMatrix W, final SparseMatrix H) {
      double cost = 0;
      try {
         cost = V.minus(W.m_times(H)).frobeniusNorm();
      }
      catch (final Exception e) {
         System.out.println(e.toString());
         System.exit(-1);
      }
      
      return cost;
   }
   
   //TODO: Figure out an acceptable range of error.
   public static double getAcceptableError() {
      return 10;
   }
   
   public static SparseMatrix updateW(final SparseMatrix V, final SparseMatrix W, final SparseMatrix H) {
      SparseMatrix RET = new SparseMatrix(1, 1);
      
      try {
         final SparseMatrix A = V.m_times(H.transpose());
         final SparseMatrix B = W.m_times(H).m_times(H.transpose());
         
         RET = W.times(A.divide(B));
         
      }
      catch (final Exception e) {
         System.out.println(e.toString());
         System.exit(-1);
      }
      
      return RET;
   }
   
   public static SparseMatrix updateH(final SparseMatrix V, final SparseMatrix W, final SparseMatrix H)
   {
      SparseMatrix RET = new SparseMatrix(1, 1);
      
      try {
         final SparseMatrix A = W.transpose().m_times(V);
         final SparseMatrix B = W.transpose().m_times(W).m_times(H);
         
         RET = H.times(A.divide(B));
         
      }
      catch (final Exception e) {
         System.out.println(e.toString());
         System.exit(-1);
      }
      
      return RET;
   }
   
   public static SimpleMatrix factor(final SimpleMatrix V) {
      final SparseMatrix W = factor(convert_matrix_to_sparse(V));
      return convert_matrix_to_simple(W);
   }
   
   public static SparseMatrix factor(final SparseMatrix V) {
      final NMFAlgorithm.Results results = getRValue(V);
      System.out.println(results.error);
      return results.W;
   }
   
   public static NMFAlgorithm.Results factor(final SparseMatrix V, int r) {
      int numIters = getNumIterations(V);
      
      if (r == 0) {
         r = 1;
      }
      
      SparseMatrix W = new SparseMatrix(V.num_rows, r);
      SparseMatrix H = new SparseMatrix(r, V.num_cols);
      
      W.randomize();
      H.randomize();
      
      do {
         H = updateH(V, W, H);
         W = updateW(V, W, H);
         System.out.println(numIters + " " + FrobeniusCost(V, W, H));
      } while (--numIters > 0);
      
      return new Results(H, FrobeniusCost(V, W, H));
   }
   
   public static class Results {
      public Results(final SparseMatrix W, final double error) {
         this.W = W;
         this.error = error;
      }
      
      public SparseMatrix W;
      public double error;
   }
   
   public static SparseMatrix convert_matrix_to_sparse(final SimpleMatrix A) {
      final SparseMatrix B = new SparseMatrix(A.numRows(), A.numCols());
      
      for (int i = 0; i < A.numRows(); i++) {
         for (int j = 0; j < A.numCols(); j++) {
            final double element = A.get(i, j);
            if (element != 0) {
               B.set_element(i, j, element);
            }
         }
      }
      
      return B;
   }
   
   public static SimpleMatrix convert_matrix_to_simple(final SparseMatrix A) {
      final SimpleMatrix C = new SimpleMatrix(A.num_rows, A.num_cols);
      
      for (int i = 0; i < A.coordinates.size(); i++) {
         C.set(A.coordinates.get(i).row, A.coordinates.get(i).col, A.coordinates.get(i).element);
      }
      
      return C;
   }
}
